"""
    Title

    author: wxz
    date: 2021-12-12
    github: https://github.com/xinzwang
"""

import numpy as np


def sim_ann(loss, dim=1, mu=0, sigma=1):
    """
    1. 初始化参数
    2. 计算loss
        2.1 比当前好，接受
        2.2 没有当前好，以概率p接受
    3. 重复步骤2若干次
    4. 减小p，温度降低
    5. 返回步骤2继续
    """

    # 0.初始化 算法超参数
    T = 5000  # 最大温度降低迭代次数
    N = 300  # 同温度下搜索次数
    r = 0.99  # 降温因子
    t = 0.5  # 初始温度

    # 1.初始化 变量
    best_x = None  # 最优解
    best_loss = 1000  # 最优loss
    x = np.random.normal(mu, sigma, [1,dim])  # 随机生成解
    cur_loss = loss(x)

    # 2.降温大循环
    epoch = 0
    while epoch < T and t > 1e-5:
        epoch += 1
        # 3.同温搜索循环
        for j in range(N):
            dx = np.random.normal(mu, sigma, dim)
            new_loss = loss(x + dx)
            # 4. 更新旧的解
            if new_loss < cur_loss:  # 好的结果，接受
                x += dx
                cur_loss = new_loss
                if cur_loss < best_loss:
                    best_x = x.copy()
                    best_loss = cur_loss
            else:  # 不好的结果
                k = np.random.random()
                dloss = new_loss - cur_loss
                if k < np.exp(-dloss / t):  # 以概率接受
                    x += dx
                    cur_loss = new_loss
                else:  # 不接受
                    pass
            pass
        t = t * r
        print('epoch=%d  t=%.5f  loss=%f  x=' % (epoch, t, cur_loss), x)
    pass
    print('Best: loss=%f x=' % (best_loss), best_x)
    return best_loss, best_x


if __name__ == '__main__':
    from funcs import *

    f = F_Rosenbrock(n=2)

    loss = f.forward

    _, _ = sim_ann(loss, dim=5)

    a = np.ones([1,5])
    print(loss(a))


